#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : 应用案例.py
# @Author: HCY
# @Date  : 2025/5/17
# @Desc  : 讯飞星火Embedding服务应用案例

import os
import time
import json
import numpy as np
import gradio as gr
from typing import List, Dict, Tuple
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

# SDK引入模型
from dwspark.config import Config
from dwspark.models import LLMEmbedding
# 日志
from loguru import logger

# 加载系统环境变量：SPARKAI_UID、SPARKAI_APP_ID、SPARKAI_API_KEY、SPARKAI_API_SECRET
config = Config()
# 创建星火Embedding实例
model = LLMEmbedding(config)

# 存储文档和向量的全局变量
documents = []
document_vectors = []

def get_embedding(text: str) -> List[float]:
    """
    获取文本的向量表示
    :param text: 输入文本
    :return: 向量表示
    """
    try:
        vector = model.get_embedding(text)
        return vector
    except Exception as e:
        logger.error(f"获取向量失败: {e}")
        return None

def calculate_similarity(text1: str, text2: str) -> float:
    """
    计算两段文本的相似度
    :param text1: 第一段文本
    :param text2: 第二段文本
    :return: 相似度得分
    """
    try:
        similarity = model.calculate_similarity(text1, text2)
        return similarity
    except Exception as e:
        logger.error(f"计算相似度失败: {e}")
        return None

def compare_texts(text1: str, text2: str):
    """
    比较两段文本的相似度
    :param text1: 第一段文本
    :param text2: 第二段文本
    :return: 相似度结果
    """
    if not text1 or not text2:
        return "请输入两段文本"
    
    try:
        similarity = calculate_similarity(text1, text2)
        if similarity is None:
            return "计算相似度失败"
        
        # 格式化相似度结果
        similarity_percentage = similarity * 100
        result = f"相似度: {similarity_percentage:.2f}%\n\n"
        
        # 根据相似度给出解释
        if similarity > 0.9:
            result += "解释: 这两段文本几乎相同或高度相似"
        elif similarity > 0.7:
            result += "解释: 这两段文本非常相似，可能表达相同的意思"
        elif similarity > 0.5:
            result += "解释: 这两段文本有一定相似性，部分内容相关"
        elif similarity > 0.3:
            result += "解释: 这两段文本相似度较低，可能有少量相关内容"
        else:
            result += "解释: 这两段文本几乎不相关"
            
        return result
    except Exception as e:
        logger.error(f"比较文本失败: {e}")
        return f"比较文本失败: {str(e)}"

def add_document(doc_title: str, doc_content: str):
    """
    添加文档到向量数据库
    :param doc_title: 文档标题
    :param doc_content: 文档内容
    :return: 添加结果
    """
    global documents, document_vectors
    
    if not doc_title or not doc_content:
        return "请输入文档标题和内容", f"当前文档数: {len(documents)}篇"
    
    try:
        # 获取文档的向量表示
        vector = get_embedding(doc_content)
        if vector is None:
            return "获取文档向量失败", f"当前文档数: {len(documents)}篇"
        
        # 添加文档和向量
        documents.append({"title": doc_title, "content": doc_content})
        document_vectors.append(vector)
        
        return f"成功添加文档: {doc_title}", f"当前文档数: {len(documents)}篇"
    except Exception as e:
        logger.error(f"添加文档失败: {e}")
        return f"添加文档失败: {str(e)}", f"当前文档数: {len(documents)}篇"

def search_documents(query: str, top_k: int = 3):
    """
    搜索相关文档
    :param query: 查询文本
    :param top_k: 返回结果数量
    :return: 搜索结果
    """
    global documents, document_vectors
    
    if not query:
        return "请输入搜索内容"
    
    if not documents:
        return "文档库为空，请先添加文档"
    
    try:
        # 获取查询的向量表示
        query_vector = get_embedding(query)
        if query_vector is None:
            return "获取查询向量失败"
        
        # 计算查询与所有文档的相似度
        similarities = []
        for i, doc_vector in enumerate(document_vectors):
            # 使用余弦相似度计算
            sim = cosine_similarity([query_vector], [doc_vector])[0][0]
            similarities.append((i, sim))
        
        # 按相似度排序
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        # 取top_k个结果
        result = "搜索结果:\n\n"
        for i in range(min(top_k, len(similarities))):
            doc_idx, sim = similarities[i]
            doc = documents[doc_idx]
            result += f"文档 {i+1}: {doc['title']}\n"
            result += f"相似度: {sim*100:.2f}%\n"
            
            # 截取部分内容显示
            content_preview = doc['content']
            if len(content_preview) > 100:
                content_preview = content_preview[:100] + "..."
            result += f"内容预览: {content_preview}\n\n"
        
        return result
    except Exception as e:
        logger.error(f"搜索文档失败: {e}")
        return f"搜索文档失败: {str(e)}"

def list_documents():
    """
    列出所有文档
    :return: 文档列表
    """
    global documents
    
    if not documents:
        return "文档库为空，请先添加文档"
    
    try:
        result = "文档列表:\n\n"
        for i, doc in enumerate(documents):
            result += f"{i+1}. {doc['title']}\n"
        
        return result
    except Exception as e:
        logger.error(f"列出文档失败: {e}")
        return f"列出文档失败: {str(e)}"

def add_sample_documents():
    """
    添加示例文档
    :return: 添加结果
    """
    sample_docs = [
        {
            "title": "人工智能简介",
            "content": "人工智能(AI)是计算机科学的一个分支，致力于创造能够模拟人类智能的机器。它包括机器学习、深度学习、自然语言处理等多个领域。现代AI技术已经在医疗、金融、交通等多个行业得到广泛应用。"
        },
        {
            "title": "机器学习基础",
            "content": "机器学习是人工智能的核心技术之一，它使计算机能够从数据中学习并做出预测或决策，而无需显式编程。常见的机器学习算法包括线性回归、决策树、随机森林、支持向量机和神经网络等。"
        },
        {
            "title": "深度学习技术",
            "content": "深度学习是机器学习的一个子领域，它使用多层神经网络来模拟人脑的工作方式。深度学习在图像识别、语音识别和自然语言处理等任务上取得了突破性进展，推动了人工智能技术的快速发展。"
        },
        {
            "title": "自然语言处理",
            "content": "自然语言处理(NLP)是人工智能的一个分支，专注于使计算机理解、解释和生成人类语言。NLP技术包括文本分类、情感分析、机器翻译、问答系统和对话系统等，广泛应用于搜索引擎、智能助手和客服机器人等场景。"
        },
        {
            "title": "计算机视觉",
            "content": "计算机视觉是人工智能的一个领域，专注于使计算机能够从图像或视频中获取高级理解。它包括图像分类、目标检测、图像分割和人脸识别等技术，应用于自动驾驶、安防监控、医疗诊断等多个领域。"
        }
    ]
    
    success_count = 0
    for doc in sample_docs:
        try:
            result, _ = add_document(doc["title"], doc["content"])
            if "成功" in result:
                success_count += 1
        except Exception as e:
            logger.error(f"添加示例文档失败: {e}")
    
    return f"成功添加 {success_count} 个示例文档", f"当前文档数: {len(documents)}篇"

def clear_documents():
    """
    清空文档库
    :return: 清空结果
    """
    global documents, document_vectors
    
    try:
        count = len(documents)
        documents = []
        document_vectors = []
        
        return f"已清空文档库，共删除 {count} 篇文档", "当前文档数: 0篇"
    except Exception as e:
        logger.error(f"清空文档库失败: {e}")
        return f"清空文档库失败: {str(e)}", f"当前文档数: {len(documents)}篇"

def cluster_documents(num_clusters: int = 3):
    """
    对文档进行聚类
    :param num_clusters: 聚类数量
    :return: 聚类结果
    """
    global documents, document_vectors
    
    if not documents:
        return "文档库为空，请先添加文档"
    
    if len(documents) < num_clusters:
        return f"文档数量({len(documents)})小于聚类数量({num_clusters})，请减少聚类数量或增加文档"
    
    try:
        from sklearn.cluster import KMeans
        
        # 将向量转换为numpy数组
        vectors_array = np.array(document_vectors)
        
        # 使用K-means聚类
        kmeans = KMeans(n_clusters=num_clusters, random_state=0).fit(vectors_array)
        labels = kmeans.labels_
        
        # 整理聚类结果
        clusters = {}
        for i, label in enumerate(labels):
            if label not in clusters:
                clusters[label] = []
            clusters[label].append(i)
        
        # 生成结果报告
        result = f"文档聚类结果 (共{num_clusters}个聚类):\n\n"
        for cluster_id, doc_indices in clusters.items():
            result += f"聚类 {cluster_id+1} (包含 {len(doc_indices)} 篇文档):\n"
            for idx in doc_indices:
                result += f"  - {documents[idx]['title']}\n"
            result += "\n"
        
        return result
    except Exception as e:
        logger.error(f"聚类文档失败: {e}")
        return f"聚类文档失败: {str(e)}"

# 构建Gradio界面
with gr.Blocks(title="星火Embedding应用") as demo:
    gr.Markdown("# 讯飞星火Embedding应用")
    
    # 文档计数
    doc_count = gr.Textbox(label="文档统计", value="当前文档数: 0篇", interactive=False)
    
    with gr.Tab("文本相似度比较"):
        gr.Markdown("## 文本相似度比较")
        with gr.Row():
            with gr.Column():
                text1 = gr.Textbox(label="文本1", placeholder="请输入第一段文本", lines=5)
                text2 = gr.Textbox(label="文本2", placeholder="请输入第二段文本", lines=5)
                compare_btn = gr.Button("比较相似度")
                similarity_result = gr.Textbox(label="比较结果", interactive=False, lines=5)
        
        # 示例文本对
        gr.Examples([
            ["今天天气真不错，阳光明媚", "今日的天气非常好，阳光充足"],
            ["人工智能技术正在快速发展", "AI技术的进步日新月异"],
            ["我喜欢吃苹果", "这台电脑是苹果品牌的"]
        ], [text1, text2])
        
        # 绑定比较按钮
        compare_btn.click(
            fn=compare_texts,
            inputs=[text1, text2],
            outputs=[similarity_result]
        )
    
    with gr.Tab("文档管理"):
        gr.Markdown("## 文档管理")
        with gr.Row():
            with gr.Column():
                doc_title = gr.Textbox(label="文档标题", placeholder="请输入文档标题")
                doc_content = gr.Textbox(label="文档内容", placeholder="请输入文档内容", lines=10)
                add_doc_btn = gr.Button("添加文档")
                add_result = gr.Textbox(label="添加结果", interactive=False)
                
                # 绑定添加文档按钮
                add_doc_btn.click(
                    fn=add_document,
                    inputs=[doc_title, doc_content],
                    outputs=[add_result, doc_count]
                )
            
            with gr.Column():
                gr.Markdown("### 文档操作")
                add_samples_btn = gr.Button("添加示例文档")
                list_docs_btn = gr.Button("列出所有文档")
                clear_docs_btn = gr.Button("清空文档库")
                docs_list = gr.Textbox(label="文档列表", interactive=False, lines=15)
                
                # 绑定文档操作按钮
                add_samples_btn.click(
                    fn=add_sample_documents,
                    inputs=[],
                    outputs=[docs_list, doc_count]
                )
                
                list_docs_btn.click(
                    fn=list_documents,
                    inputs=[],
                    outputs=[docs_list]
                )
                
                clear_docs_btn.click(
                    fn=clear_documents,
                    inputs=[],
                    outputs=[docs_list, doc_count]
                )
    
    with gr.Tab("语义搜索"):
        gr.Markdown("## 语义搜索")
        with gr.Row():
            with gr.Column():
                search_query = gr.Textbox(label="搜索内容", placeholder="请输入搜索内容")
                top_k = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="返回结果数量")
                search_btn = gr.Button("搜索")
                search_result = gr.Textbox(label="搜索结果", interactive=False, lines=15)
        
        # 示例搜索查询
        gr.Examples([
            "什么是人工智能?",
            "机器学习有哪些常见算法?",
            "深度学习和机器学习有什么区别?",
            "自然语言处理的应用场景"
        ], search_query)
        
        # 绑定搜索按钮
        search_btn.click(
            fn=search_documents,
            inputs=[search_query, top_k],
            outputs=[search_result]
        )
    
    with gr.Tab("文档聚类"):
        gr.Markdown("## 文档聚类分析")
        with gr.Row():
            with gr.Column():
                num_clusters = gr.Slider(minimum=2, maximum=10, value=3, step=1, label="聚类数量")
                cluster_btn = gr.Button("进行聚类")
                cluster_result = gr.Textbox(label="聚类结果", interactive=False, lines=15)
        
        # 绑定聚类按钮
        cluster_btn.click(
            fn=cluster_documents,
            inputs=[num_clusters],
            outputs=[cluster_result]
        )

if __name__ == "__main__":
    demo.queue().launch()